#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Kernel command line handler class
"""

import os
import subprocess
import logging
from typing import List

logger = logging.getLogger(__name__)

class KernelCmdline:
    """Kernel command line handler class"""
    
    def __init__(self, image_path: str):
        """
        Initialize
        
        Args:
            image_path: image path
        """
        self.image_path = image_path
        self._check_dependencies()
    
    def _check_dependencies(self):
        """Check if required tools are installed"""
        try:
            subprocess.run(['virt-customize', '--version'], 
                         stdout=subprocess.PIPE, 
                         stderr=subprocess.PIPE, 
                         check=True)
            subprocess.run(['guestfish', '--version'],
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE,
                         check=True)
            logger.debug("Required tools found")
        except (subprocess.SubprocessError, FileNotFoundError):
            raise RuntimeError("Required tools not found, please install libguestfs-tools package first")
    
    def get_current_cmdline(self) -> str:
        """
        Get current kernel command line parameters
        
        Returns:
            str: Current kernel command line parameters
        """
        try:
            # Use guestfish to read grub.cfg file
            cmd = [
                'guestfish', '--ro', '-a', self.image_path,
                '-i',  # 启用自动挂载
                'cat', '/boot/grub/grub.cfg'
            ]
            
            result = subprocess.run(cmd,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 text=True,
                                 check=True)
            
            # Parse grub.cfg content, find the first linux command line
            cmdline = ""
            for line in result.stdout.splitlines():
                line = line.strip()
                if 'linux' in line and '/vmlinuz' in line:
                    # Extract kernel command line parameters
                    parts = line.split()
                    # Skip 'linux' command and kernel path
                    if len(parts) > 2:
                        cmdline = ' '.join(parts[2:])
                        break
            
            if not cmdline:
                logger.warning("Failed to get cmdline from grub.cfg")
                return ""

            logger.debug(f"Current cmdline: {cmdline}")
            return cmdline
            
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to get current cmdline: {e.stderr}")
            raise RuntimeError(f"Failed to get current cmdline: {e.stderr}")
    
    def replace_cmdline(self, new_cmdline: str) -> bool:
        """
        Completely replace kernel command line
        
        Args:
            new_cmdline: New full command line string
        
        Returns:
            bool: Success or not
        """
        try:
            # Read current grub.cfg content
            cmd = ['guestfish', '--ro', '-a', self.image_path,
                  '-i', 'cat', '/boot/grub/grub.cfg']
            result = subprocess.run(cmd,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.PIPE,
                                 text=True,
                                 check=True)
            
            grub_content = result.stdout
            new_content = []
            in_menuentry = False
            
            # Process each line
            for line in grub_content.splitlines():
                if line.startswith('menuentry '):
                    in_menuentry = True
                elif line.startswith('}'):
                    in_menuentry = False
                
                if in_menuentry and ('linux' in line and '/vmlinuz' in line):
                    # Found the kernel command line
                    parts = line.split(None, 2)
                    if len(parts) >= 2:
                        # Keep command and kernel path, replace all parameters
                        new_line = f"{parts[0]} {parts[1]} {new_cmdline}"
                        new_content.append(new_line)
                        continue
                
                new_content.append(line)
            
            # Write new content to temp file
            temp_file = '/tmp/new_grub.cfg'
            write_cmd = ['guestfish', '-a', self.image_path, '-i',
                        'write', '/boot/grub/grub.cfg', '\n'.join(new_content)]
            
            subprocess.run(write_cmd,
                         stdout=subprocess.PIPE,
                         stderr=subprocess.PIPE,
                         text=True,
                         check=True)
            
            logger.info(f"Kernel command line updated to: {new_cmdline}")
            return True
            
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to replace cmdline: {e.stderr}")
            return False
    
    def update_cmdline(self, params: List[str]) -> bool:
        """
        Update specific kernel command line parameters
        
        Args:
            params: List of parameters to update
        
        Returns:
            bool: Success or not
        """
        try:
            # Get current cmdline
            current_cmdline = self.get_current_cmdline()
            
            # Define unique parameters
            UNIQUE_PARAMS = {
                'root',      # root filesystem
                'init',      # init process
                'resume',    # swap partition
                'rootflags', # root filesystem mount options
                'rootfstype',# root filesystem type
                'ro',        # read-only flag
                'rw'        # read-write flag
            }
            
            # Use dict to store parameters, key is param name, value is list of values
            param_dict = {}
            
            # Process existing parameters
            for param in current_cmdline.split():
                key = param
                value = None
                if '=' in param:
                    key, value = param.split('=', 1)
                
                if key not in param_dict:
                    param_dict[key] = []
                
                if key in UNIQUE_PARAMS:
                    # Unique parameter, clear old value
                    param_dict[key] = []
                
                # Add to list if value not duplicate
                if value is not None:
                    if value not in param_dict[key]:
                        param_dict[key].append(value)
                elif not param_dict[key]:  # Only keep once for valueless param
                    param_dict[key].append(None)
            
            # Process new parameters
            for param in params:
                key = param
                value = None
                if '=' in param:
                    key, value = param.split('=', 1)
                
                if key not in param_dict:
                    param_dict[key] = []
                
                if key in UNIQUE_PARAMS:
                    # Unique parameter, clear old value
                    param_dict[key] = []
                
                # Add to list if value not duplicate
                if value is not None:
                    if value not in param_dict[key]:
                        param_dict[key].append(value)
                elif not param_dict[key]:  # Only keep once for valueless param
                    param_dict[key].append(None)
            
            # Assemble final cmdline
            final_cmdline = []
            
            # Process all parameters
            for key, values in param_dict.items():
                if key in UNIQUE_PARAMS:
                    # Only use last value for unique param
                    if values:
                        value = values[-1]
                        if value is not None:
                            final_cmdline.append(f"{key}={value}")
                        else:
                            final_cmdline.append(key)
                else:
                    # For non-unique params, add all non-duplicate values
                    for value in values:
                        if value is not None:
                            final_cmdline.append(f"{key}={value}")
                        else:
                            final_cmdline.append(key)
            
            # Generate final cmdline string
            new_cmdline = ' '.join(final_cmdline)
            logger.debug(f"Updated cmdline: {new_cmdline}")
            
            # Call replace_cmdline to update
            return self.replace_cmdline(new_cmdline)
            
        except Exception as e:
            logger.error(f"Failed to update cmdline parameters: {e}")
            return False 